/*
 * Copyright (c) 2023 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ecmascript/dfx/hprof/heap_sampling.h"

#include <ucontext.h>
#if defined(PANDA_TARGET_ARM64)
#include <sys/ucontext.h>
#endif
#include <unistd.h>

#include "ecmascript/compiler/assembler/assembler.h"
#include "ecmascript/frames.h"
#include "ecmascript/aot_file_manager.h"
#include "ecmascript/jspandafile/js_pandafile_manager.h"

namespace panda::ecmascript {
HeapSampling::HeapSampling(const EcmaVM *vm, Heap *const heap, uint64_t interval, int stackDepth)
    : vm_(vm),
      heap_(heap),
      rate_(interval),
      stackDepth_(stackDepth),
      allocationInspector_(heap, interval, this)
{
    samplingInfo_ = std::make_unique<struct SamplingInfo>();
    samplingInfo_->head.id = 1;
    samplingInfo_->head.callFrameInfo.functionName = "(root)";
    heap_->AddAllocationInspectorToAllSpaces(&allocationInspector_);
}

HeapSampling::~HeapSampling()
{
    heap_->ClearAllocationInspectorFromAllSpaces();
}

std::unique_ptr<struct SamplingInfo> HeapSampling::GetAllocationProfile()
{
    return std::move(samplingInfo_);
}

void HeapSampling::ImplementSampling([[maybe_unused]]Address addr, [[maybe_unused]]size_t size)
{
    GetStack();
    AllocationNode *node = PushAndGetNode();
    node->allocations[size]++;
    node->selfSize += size;
    samplingInfo_->samples.push_back(Sample(size, node->id, CreateSampleId()));
}

bool HeapSampling::PushStackInfo(struct MethodKey &methodKey)
{
    if (UNLIKELY(frameStack_.size() >= static_cast<size_t>(stackDepth_))) {
        return false;
    }
    frameStack_.push_back(methodKey);
    return true;
}

bool HeapSampling::PushFrameInfo(const FrameInfoTemp &frameInfoTemp)
{
    if (UNLIKELY(frameInfoTemps_.size() >= static_cast<size_t>(stackDepth_))) {
        return false;
    }
    frameInfoTemps_.push_back(frameInfoTemp);
    return true;
}

void HeapSampling::ResetFrameLength()
{
    stackInfoMap_.clear();
    frameStack_.clear();
}

std::string HeapSampling::GetRunningState(RunningState state, kungfu::DeoptType type)
{
    std::string result;
    if (state == RunningState::AOT && type != kungfu::DeoptType::NOTCHECK) {
        state = RunningState::AINT;
    }
    switch (state) {
        case RunningState::GC:
            result = "(GC)";
            break;
        case RunningState::CINT:
            result = "(CINT)";
            break;
        case RunningState::AINT:
            result = "(AINT)";
            break;
        case RunningState::AOT:
            result = "(AOT)";
            break;
        case RunningState::BUILTIN:
            result = "(BUILTIN)";
            break;
        case RunningState::NAPI:
            result = "(NAPI)";
            break;
        case RunningState::ARKUI_ENGINE:
            result = "(ARKUI_ENGINE)";
            break;
        case RunningState::RUNTIME:
            result = "(RUNTIME)";
            break;
        default:
            result = "(OTHER)";
            break;
    }
    return result;
}


void HeapSampling::GetStack()
{
    ResetFrameLength();
    JSThread *thread = vm_->GetAssociatedJSThread();
    JSTaggedType *frame = const_cast<JSTaggedType *>(thread->GetCurrentFrame());
    if (frame != nullptr) {
        FrameIterator it(frame, thread);
        for (; !it.Done(); it.Advance<>()) {
            auto method = it.CheckAndGetMethod();
            if (method == nullptr) {
                continue;
            }
            const JSPandaFile *jsPandaFile = method->GetJSPandaFile();
            struct MethodKey methodKey;
            void *methodIdentifier = JsStackGetter::GetMethodIdentifier(method, it);
            if (methodIdentifier == nullptr) {
                return;
            }
            methodKey.methodIdentifier = methodIdentifier;
            if (stackInfoMap_.count(methodKey) == 0) {
                struct FrameInfoTemp codeEntry;
                if (UNLIKELY(!JsStackGetter::ParseMethodInfo(methodKey, it, jsPandaFile, vm_, codeEntry) &&
                    !PushFrameInfo(codeEntry))) {
                    return;
                }
            }
            if (UNLIKELY(!PushStackInfo(methodKey))) {
                return;
            }
        }
    }
}
void HeapSampling::FillScriptIdAndStore()
{
    size_t len = frameInfoTemps_.size();
    if (len == 0) {
        return;
    }
    struct CallFrameInfo callframeInfo;
    for (size_t i = 0; i < len; ++i) {
        callframeInfo.url = frameInfoTemps_[i].url;
        auto iter = scriptIdMap_.find(callframeInfo.url);
        if (iter == scriptIdMap_.end()) {
            scriptIdMap_.emplace(callframeInfo.url, scriptIdMap_.size() + 1);
            callframeInfo.scriptId = static_cast<int>(scriptIdMap_.size());
        } else {
            callframeInfo.scriptId = iter->second;
        }
        callframeInfo.functionName = frameInfoTemps_[i].functionName;
        callframeInfo.codeType = frameInfoTemps_[i].codeType;
        callframeInfo.columnNumber = frameInfoTemps_[i].columnNumber;
        callframeInfo.lineNumber = frameInfoTemps_[i].lineNumber;
        stackInfoMap_.emplace(frameInfoTemps_[i].methodKey, callframeInfo);
    }
    frameInfoTemps_.clear();
}

struct AllocationNode *HeapSampling::PushAndGetNode()
{
    FillScriptIdAndStore();
    struct AllocationNode *node = &samplingInfo_->head;
    int frameLen = static_cast<int>(frameStack_.size()) - 1;
    for (; frameLen >= 0; frameLen--) {
        CallFrameInfo frameInfo = GetMethodInfo(frameStack_[frameLen]);
        node = FindOrAddNode(node, frameStack_[frameLen]);
    }
    return node;
}

struct CallFrameInfo HeapSampling::GetMethodInfo(MethodKey methodKey)
{
    struct CallFrameInfo frameInfo;
    auto iter = stackInfoMap_.find(methodKey);
    if (iter != stackInfoMap_.end()) {
        frameInfo = iter->second;
    }
    return frameInfo;
}

struct AllocationNode *HeapSampling::FindOrAddNode(struct AllocationNode *node, MethodKey &methodKey)
{
    struct AllocationNode *childNode = nullptr;
    if (node->children.count(methodKey) != 0) {
        childNode = &(node->children[methodKey]);
    }
    if (childNode == nullptr) {
        struct AllocationNode tempNode;
        tempNode.callFrameInfo = GetMethodInfo(methodKey);
        tempNode.id = CreateNodeId();
        node->children[methodKey] = tempNode;
    }
    return childNode;
}

uint32_t HeapSampling::CreateNodeId()
{
    return ++nodeId_;
}

uint64_t HeapSampling::CreateSampleId()
{
    return ++SampleId_;
}
}  // namespace panda::ecmascript
